

using namespace AscendC;



template <typename T>
class TokenShiftB1Align
{
public:
    __aicore__ inline TokenShiftB1Align() {}
    __aicore__ inline void Init(
        GM_ADDR x, 
        GM_ADDR rwkvag, 
        GM_ADDR h0, 
        GM_ADDR xr, 
        GM_ADDR xw, 
        GM_ADDR xk,
        GM_ADDR xv, 
        GM_ADDR xa, 
        GM_ADDR xg, 
        GM_ADDR ht, 
        GM_ADDR workspace,
        uint32_t aivNum,
        uint32_t B_length,
        uint32_t T_length,
        uint32_t C_length,
        uint32_t tiling_size,
        uint32_t core_rows,
        uint32_t core_rows_remain,
        uint32_t X_length_align,
        TPipe* pipeIn)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero");
        this->pipe = pipeIn;
        this->aivNum = aivNum;
        this->startPointer = core_rows * C_length * GetBlockIdx();
        this->core_rows = core_rows + (GetBlockNum() == GetBlockIdx() + 1 ? core_rows_remain : 0);
        this->B_length = B_length;
        this->T_length = T_length;
        this->C_length = C_length;
        this->tiling_size = tiling_size;
        this->ALIGN_NUM = 32 / sizeof(T);
        uint32_t core_length = this->core_rows * C_length;
        this->prevLinePointer = core_rows * C_length * GetBlockIdx() - C_length;
        // printf("prevLinePointer: %u\n", this->prevLinePointer);

        if(GetBlockIdx() == 0) {
            x_global.SetGlobalBuffer((__gm__ T *)x, core_length);
        } else {
            x_global.SetGlobalBuffer((__gm__ T *)x + this->startPointer - C_length, core_length + C_length);
        }
        rwkvag_global.SetGlobalBuffer((__gm__ T *)rwkvag, 6 * C_length);
        h0_global.SetGlobalBuffer((__gm__ T *)h0 , C_length);
        
        xr_global.SetGlobalBuffer((__gm__ T *)xr + this->startPointer, core_length);   
        xw_global.SetGlobalBuffer((__gm__ T *)xw + this->startPointer, core_length);   
        xk_global.SetGlobalBuffer((__gm__ T *)xk + this->startPointer, core_length);   
        xv_global.SetGlobalBuffer((__gm__ T *)xv + this->startPointer, core_length);   
        xa_global.SetGlobalBuffer((__gm__ T *)xa + this->startPointer, core_length);   
        xg_global.SetGlobalBuffer((__gm__ T *)xg + this->startPointer, core_length);        
        ht_global.SetGlobalBuffer((__gm__ T *)ht, C_length); 
        // pipe->InitBuffer(inQueueXBind, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueRwkvag, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueH0, BUFFER_NUM, this->tiling_size * sizeof(T));
        // pipe->InitBuffer(inQueueXNextLine, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueXNextLine, 2, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXR, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXW, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXK, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXV, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXA, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXG, BUFFER_NUM, this->tiling_size * sizeof(T));

        pipe->InitBuffer(QueueTmpX, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp0, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp1, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp2, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp3, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp4, this->tiling_size * sizeof(T));
        pipe->InitBuffer(QueueTmp5, this->tiling_size * sizeof(T));

        g_xLocal = QueueTmpX.Get<T>();
        g_RWKVAG0Local = QueueTmp0.Get<T>();
        g_RWKVAG1Local = QueueTmp1.Get<T>();
        g_RWKVAG2Local = QueueTmp2.Get<T>();
        g_RWKVAG3Local = QueueTmp3.Get<T>();
        g_RWKVAG4Local = QueueTmp4.Get<T>();
        g_RWKVAG5Local = QueueTmp5.Get<T>();

        DataCopy(g_RWKVAG0Local, rwkvag_global, C_length);   
        DataCopy(g_RWKVAG1Local, rwkvag_global[C_length], C_length); 
        DataCopy(g_RWKVAG2Local, rwkvag_global[C_length * 2], C_length);
        DataCopy(g_RWKVAG3Local, rwkvag_global[C_length * 3], C_length);
        DataCopy(g_RWKVAG4Local, rwkvag_global[C_length * 4], C_length);
        DataCopy(g_RWKVAG5Local, rwkvag_global[C_length * 5], C_length);

        PipeBarrier<PIPE_ALL>(); 

    }

    __aicore__ inline void Process() {
        // int32_t loopcount = this->C_length / this->tiling_size;

        // for(int32_t p = 0; p < loopcount; p++) {
        //     if(GetBlockIdx() == 0) {
        //         CopyInH0(p * this->tiling_size, this->tiling_size);
        //         for(int32_t k = 0; k < this->core_rows; k++) {
        //             CopyInXNextLine(k * this->C_length + p * this->tiling_size, this->tiling_size);
        //             Compute(p * this->tiling_size, this->tiling_size);
        //             CopyOutX(k * this->C_length + p * this->tiling_size, this->tiling_size);
        //         }
        //     } else {
        //         CopyInXPrevLine(p * this->tiling_size, this->tiling_size);
        //         for(int32_t k = 0; k < this->core_rows; k++) {
        //             CopyInXNextLine(this->C_length + k * this->C_length + p * this->tiling_size, this->tiling_size);
        //             Compute(p * this->tiling_size, this->tiling_size);
        //             CopyOutX(this->C_length + k * this->C_length + p * this->tiling_size, this->tiling_size);
        //         }
        //     }
        //     if(GetBlockNum() == GetBlockIdx() + 1) {
        //         CopyOutHT(p * this->tiling_size, this->tiling_size);
        //     }
        // }

        if(GetBlockIdx() == 0) {
            CopyInH0(0, this->tiling_size);
            for(int32_t k = 0; k < this->core_rows; k++) {
                CopyInXNextLine(k * this->C_length, this->tiling_size);
                Compute(0, this->tiling_size);
                CopyOutX(k * this->C_length, this->tiling_size);
            }
        } else {
            CopyInXPrevLine(0, this->tiling_size);
            for(int32_t k = 0; k < this->core_rows; k++) {
                CopyInXNextLine(this->C_length + k * this->C_length, this->tiling_size);
                Compute(0, this->tiling_size);
                CopyOutX(k * this->C_length, this->tiling_size);
            }
        }
        SyncAll<true>();
        uint32_t blockidx = GetBlockIdx();
        // printf("blockidx:%u\n", blockidx);
        if(this->aivNum == GetBlockIdx() + 1) {
            CopyOutHT(0, this->tiling_size);
        }
               
    }
    
    __aicore__ inline void Compute(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xNextLineLocal = inQueueXNextLine.DeQue<T>();
        // LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        LocalTensor<T> XRLocal = outQueueXR.AllocTensor<T>();
        LocalTensor<T> XWLocal = outQueueXW.AllocTensor<T>();
        LocalTensor<T> XKLocal = outQueueXK.AllocTensor<T>();
        LocalTensor<T> XVLocal = outQueueXV.AllocTensor<T>();
        LocalTensor<T> XALocal = outQueueXA.AllocTensor<T>();
        LocalTensor<T> XGLocal = outQueueXG.AllocTensor<T>();

        Sub(g_xLocal, g_xLocal, xNextLineLocal, length);
        Mul(XRLocal, g_xLocal, g_RWKVAG0Local, length);
        Add(XRLocal, XRLocal, xNextLineLocal, length);
        outQueueXR.EnQue(XRLocal);  

        Mul(XWLocal, g_xLocal, g_RWKVAG1Local, length);
        Add(XWLocal, XWLocal, xNextLineLocal, length);
        outQueueXW.EnQue(XWLocal);  
        

        Mul(XKLocal, g_xLocal, g_RWKVAG2Local, length);
        Add(XKLocal, XKLocal, xNextLineLocal, length);
        outQueueXK.EnQue(XKLocal);  

        Mul(XVLocal, g_xLocal, g_RWKVAG3Local, length);
        Add(XVLocal, XVLocal, xNextLineLocal, length);
        outQueueXV.EnQue(XVLocal); 

        Mul(XALocal, g_xLocal, g_RWKVAG4Local, length);
        Add(XALocal, XALocal, xNextLineLocal, length);
        outQueueXA.EnQue(XALocal); 


        Mul(XGLocal, g_xLocal, g_RWKVAG5Local, length);
        Add(XGLocal, XGLocal, xNextLineLocal, length);
        outQueueXG.EnQue(XGLocal);      
        
        Adds(g_xLocal, xNextLineLocal, T(0), length);

        // inQueueX.EnQue(xLocal);
        inQueueXNextLine.FreeTensor(xNextLineLocal);
    }    

    __aicore__ inline void CopyInH0(uint32_t startpos, uint32_t length) {
        // LocalTensor<T> h0Local = inQueueH0.AllocTensor<T>();
        DataCopy(g_xLocal, h0_global[startpos], length);   
        PipeBarrier<PIPE_ALL>(); 
        // inQueueH0.EnQue(h0Local);         
    }       
    __aicore__ inline void CopyInXPrevLine(uint32_t startpos, uint32_t length) {
        DataCopy(g_xLocal, x_global[startpos], length);   
        PipeBarrier<PIPE_ALL>();      
    }       

    __aicore__ inline void CopyInXNextLine(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xLocal = inQueueXNextLine.AllocTensor<T>();
        DataCopy(xLocal, x_global[startpos], length);   
        inQueueXNextLine.EnQue(xLocal);  
    }

    // __aicore__ inline void CopyInRline(uint32_t startpos, uint32_t length) {
    //     LocalTensor<T> rwkvagLocal = inQueueRwkvag.AllocTensor<T>();
    //     DataCopy(rwkvagLocal, rwkvag_global[startpos], length);   
    //     inQueueRwkvag.EnQue(rwkvagLocal);  
    // }

    __aicore__ inline void CopyOutHT(uint32_t startpos, uint32_t length) {
        // LocalTensor<T> xLocal = inQueueXBind.DeQue<T>();
        // printf("CopyOutHT");
        DataCopy(ht_global[startpos], g_xLocal, length);  
        PipeBarrier<PIPE_ALL>();  
    }

    // __aicore__ inline void CopyOutHTUnalign(uint32_t startpos, uint32_t length) {
    //     uint32_t process_length = length / this->ALIGN_NUM * this->ALIGN_NUM;

    //     DataCopy(ht_global[startpos], g_xLocal, process_length); 
    //     for(int32_t m = process_length; m < length; m++) {
    //         ht_global.SetValue(m + startpos, g_xLocal.GetValue(m));
    //     }    
    //     PipeBarrier<PIPE_ALL>();

    // }

    __aicore__ inline void CopyOutX(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xrLocal = outQueueXR.DeQue<T>();
        DataCopy(xr_global[startpos], xrLocal, length); 
        outQueueXR.FreeTensor(xrLocal);

        LocalTensor<T> xwLocal = outQueueXW.DeQue<T>();
        DataCopy(xw_global[startpos], xwLocal, length); 
        outQueueXW.FreeTensor(xwLocal);

        LocalTensor<T> xkLocal = outQueueXK.DeQue<T>();
        DataCopy(xk_global[startpos], xkLocal, length); 
        outQueueXK.FreeTensor(xkLocal);

        LocalTensor<T> xvLocal = outQueueXV.DeQue<T>();
        DataCopy(xv_global[startpos], xvLocal, length); 
        outQueueXV.FreeTensor(xvLocal);

        LocalTensor<T> xaLocal = outQueueXA.DeQue<T>();
        DataCopy(xa_global[startpos], xaLocal, length); 
        outQueueXA.FreeTensor(xaLocal);

        LocalTensor<T> xgLocal = outQueueXG.DeQue<T>();
        DataCopy(xg_global[startpos], xgLocal, length); 
        outQueueXG.FreeTensor(xgLocal);
        PipeBarrier<PIPE_ALL>();
    }

private:
    GlobalTensor<T> x_global;
    GlobalTensor<T> rwkvag_global;
    GlobalTensor<T> h0_global;
    GlobalTensor<T> xr_global, xw_global, xk_global, xv_global, xa_global, xg_global, ht_global; 

    // TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> inQueueXBind;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueRwkvag, inQueueH0;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueXNextLine;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueXR, outQueueXW, outQueueXK, outQueueXV, outQueueXA, outQueueXG;

    TBuf<QuePosition::VECCALC> QueueTmpX, QueueTmp0, QueueTmp1, QueueTmp2, QueueTmp3, QueueTmp4, QueueTmp5;
    uint32_t B_length;
    uint32_t T_length;
    uint32_t C_length;
    uint32_t tiling_size;
    uint32_t core_rows;
    uint32_t core_rows_remain;
    uint32_t startPointer;
    uint32_t X_length_align;
    uint32_t ALIGN_NUM;
    uint32_t prevLinePointer;
    uint32_t aivNum;

    LocalTensor<T> g_xLocal;
    LocalTensor<T> g_RWKVAG0Local;
    LocalTensor<T> g_RWKVAG1Local;
    LocalTensor<T> g_RWKVAG2Local;
    LocalTensor<T> g_RWKVAG3Local;
    LocalTensor<T> g_RWKVAG4Local;
    LocalTensor<T> g_RWKVAG5Local;
    AscendC::TPipe* pipe;
};